import sys
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QVBoxLayout, QPushButton
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import QTimer
import cv2

class MyWindow(QWidget):
    def __init__(self):
        super().__init__()

        # 初始化 OpenCV 相关设置
        # self.cap = cv2.VideoCapture(1)
        self.cap = cv2.VideoCapture('./assets/video/2.mp4')
        self.face_detector = cv2.CascadeClassifier('./haarcascades/haarcascade_frontalface_alt2.xml')
        self.recognizer = cv2.face.LBPHFaceRecognizer.create()
        self.recognizer.read('./data/02.yml')

        self.setWindowTitle('hello world')

        self.image_label = QLabel(self)
        self.btn = QPushButton('进门')
        self.btn.setStyleSheet('height: 40')

        # 设置布局
        layout = QVBoxLayout()
        layout.addWidget(self.image_label)
        layout.addWidget(self.btn)
        self.setLayout(layout)

        # 创建定时器，每隔一段时间更新图像
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update_image)
        self.timer.start(20)  # 每20毫秒更新一次

    def update_image(self):
        # 从摄像头中读取图像
        ret, frame = self.cap.read()

        if ret:
            # OpenCV 处理图像的部分（在这里示例中仅进行灰度化处理）
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = self.face_detector.detectMultiScale(gray, 1.1, 5)
            for x, y, w, h in faces:
                tags, confidence = self.recognizer.predict(gray[y: y + h, x: x + w])
                print('标签ID：', tags, '置信度：', confidence)
                color = (0, 0, 255) if confidence > 60 else (0, 255, 0)
                cv2.rectangle(frame, (x, y), (x + w, y + h), color, thickness=2)
                cv2.putText(frame, 'ID: ' + str(tags), (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, color)

            # 将 OpenCV 处理后的图像转换为 Qt 的图像格式
            h, w, c = frame.shape
            print('w:' + str(w) + ' h:' + str(h) + ' c:' + str(c))
            bytes_per_line = 3 * w
            qt_image = QImage(frame.data, w, h, bytes_per_line, QImage.Format_BGR888)

            # 将 Qt 图像显示到 QLabel 中
            pixmap = QPixmap.fromImage(qt_image)
            self.image_label.setPixmap(pixmap)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MyWindow()
    window.show()
    sys.exit(app.exec_())